from fastapi import FastAPI
from starlette.websockets import WebSocket, WebSocketDisconnect

from mcp_server.mcp_client import MCPClient


async def init_ai_websocket(app: FastAPI):
    """
    初始化AI聊天WebSocket端点

    该函数为FastAPI应用注册一个WebSocket端点，用于处理AI聊天功能。
    每个连接的用户都会获得独立的对话上下文，并通过MCP客户端与AI服务器通信。

    参数:
        app (FastAPI): FastAPI应用实例，用于注册WebSocket路由

    返回值:
        无返回值
    """

    @app.websocket("/ws/chat")
    async def websocket_endpoint(websocket: WebSocket):
        """
        WebSocket端点处理函数

        处理客户端WebSocket连接，接收用户消息，通过MCP客户端获取AI回复，
        并将回复流式发送给客户端。

        参数:
            websocket (WebSocket): WebSocket连接对象

        返回值:
            无返回值
        """
        # 接受WebSocket连接
        await websocket.accept()

        # 为每个用户创建唯一的ID和对话上下文
        user_id = id(websocket)
        user_contexts = {}
        user_contexts[user_id] = [{"role": "system", "content": "你是一个有帮助的助手。"}]

        # 初始化MCP客户端并连接到AI服务器
        client = MCPClient()
        await client.connect_to_server('mcp_server/mcp_server.py')

        try:
            # 持续监听和处理用户消息
            while True:
                # 接收用户发送的消息
                user_msg = await websocket.receive_text()
                user_contexts[user_id].append({"role": "user", "content": user_msg})
                await websocket.send_json({"role": "user", "content": user_msg})

                # 向客户端发送开始响应的标记
                assistant_reply = ""
                response = client.put_query(user_msg)
                await websocket.send_json({"start": True})

                # 流式接收AI回复并发送给客户端
                async for content_piece in response:
                    assistant_reply += content_piece
                    await websocket.send_json({"role": "assistant", "content": content_piece})

                # 向客户端发送完成响应的标记
                await websocket.send_json({"done": True})

        except WebSocketDisconnect:
            # 处理WebSocket断开连接的情况
            print("WebSocket 断开连接")
            # 清理上下文
            user_contexts.pop(user_id, None)
            await client.cleanup()
